import pyaudio
import wave
import datetime
import numpy as np
import matplotlib.pyplot as plt
import time

# 配置音频输入
format = pyaudio.paInt16
channels = 1
rate = 44100
chunk_size = 1024
record_seconds = 2

# 创建一个PyAudio对象
audio = pyaudio.PyAudio()

# 打开音频输入流
stream = audio.open(format=format,
                    channels=channels,
                    rate=rate,
                    input=True,
                    frames_per_buffer=chunk_size)

def count_hit_from_wav(wav_file):
    # 读取音频文件数据
    with wave.open(wav_file, 'rb') as wf:
        # 获取音频参数
        channels = wf.getnchannels()
        sample_width = wf.getsampwidth()
        frame_rate = wf.getframerate()
        num_frames = wf.getnframes()

        # 读取音频数据
        audio_data = wf.readframes(num_frames)

    # 将音频数据转换为一维numpy数组
    if sample_width == 1:
        audio_data = np.frombuffer(audio_data, dtype=np.uint8)
    elif sample_width == 2:
        audio_data = np.frombuffer(audio_data, dtype=np.int16)
    elif sample_width == 4:
        audio_data = np.frombuffer(audio_data, dtype=np.int32)

    # 检查波形图峰值是否超过5000
    hit_count = np.count_nonzero(audio_data > 5000)

    return hit_count

input("按回车键开始录制音频...")

print("录音开始...")
while True:
    # 读取指定秒数的音频数据
    frames = []
    for i in range(int(rate / chunk_size * record_seconds)):
        data = stream.read(chunk_size)
        frames.append(data)

    # 将音频数据保存为WAV文件
    current_time = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"audio_{current_time}.wav"

    with wave.open(filename, "wb") as wf:
        wf.setnchannels(channels)
        wf.setsampwidth(audio.get_sample_size(format))
        wf.setframerate(rate)
        wf.writeframes(b''.join(frames))

    print(f"音频已保存为 {filename}")

    hit_count = count_hit_from_wav(filename)

    print(f"HIT: {hit_count}")  # 打印敲击次数

    # 延迟1/10秒后再次启动计数程序
    time.sleep(0.1)
    
    user_input = input("按回车键继续录制音频，输入其他字符退出程序：")
    if user_input != "":
        break

print("录音结束...")
